# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/8 上午10:40
# @Author  : shutao
# @FileName: house_dataset.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import numpy as np
import cv2 as cv
import random
import json
import os
from loguru import logger

from paddle.io import Dataset
from .base_dataset import BaseDataset
from paddle.vision import ColorJitter, RandomHorizontalFlip, RandomVerticalFlip


class HouseMapDataset(BaseDataset):
    def __init__(self, is_train=True,
                 dataset_root="E:\\work\\Dataset\\house_dataset",
                 datafile=None,
                 img_size=(400, 400),
                 label_index=None):
        super(HouseMapDataset, self).__init__()
        assert os.path.exists(dataset_root), "The specified data set root directory does not exist."
        if datafile is None:
            if is_train:
                self.datafile = os.path.join(dataset_root, "train.list")
            else:
                self.datafile = os.path.join(dataset_root, "val.list")
        else:
            self.datafile = os.path.join(dataset_root, datafile)
        assert os.path.exists(self.datafile), "The specified dataset list file does not exist."

        # 标签文件解析设置
        self.img_size = img_size
        if label_index is None:
            self.label_index = {"flat_roofs": 0, "tile_roofs": 1, "shanty_house_roofs": 2, "asbestos_tile_roofs": 3}
        else:
            self.label_index = label_index
        self.mask_shape = (len(self.label_index), img_size[0], img_size[1])

        # 数据列表配置
        self.is_train = is_train
        self.dataset_root = dataset_root
        self.datalist = []
        self.get_data_list()

        # 数据增强配置
        self.color_jitter = ColorJitter(0.4, 0.2, 0.2, 0.2)  # 随机调整图像的亮度，对比度，饱和度和色调
        self.random_horizontal_flip = RandomHorizontalFlip(1)  # 随机水平翻转
        self.random_vertical_flip = RandomVerticalFlip(1)  # 随机垂直翻转

    def get_data_list(self):
        with open(self.datafile, 'r') as f:
            datalines = f.readlines()
        for idx, _line in enumerate(datalines):
            data_sets = [
                os.path.join(self.dataset_root, file) for file in _line.replace("\n", "").split("\t")
            ]
            self.datalist.append(data_sets)

    def process(self, img, seg_mask):
        """
        数据增强
        :param seg_mask:
        :param img:
        :return:
        """
        img = self.color_jitter(img)
        if random.random() > 0.5:  # 水平翻转概率
            img = self.random_horizontal_flip(img)
            seg_mask = seg_mask[:, :, ::-1]
            # self.logger.info("水平翻转...")
        if random.random() > 0.5:  # 垂直翻转概率
            img = self.random_vertical_flip(img)
            seg_mask = seg_mask[:, ::-1, :]
            # self.logger.info("垂直翻转")
        return img, seg_mask

    def __getitem__(self, item):
        img_path, seg_mask_path = self.datalist[item][0], self.datalist[item][1]
        origin_img = cv.imread(img_path)

        # 读取分割标签图
        seg_mask = np.zeros(self.mask_shape, dtype='int32')
        with open(seg_mask_path, 'r') as f:
            seg_labels = json.load(f)  # 标签文件
        img_h, img_w = seg_labels["imageHeight"], seg_labels["imageWidth"]  # 图片高宽
        for seg_label in seg_labels["shapes"]:  # 读取所有包围框
            label = seg_label["label"]  # 标签类别
            points = [
                [loc[0] * self.img_size[0] // img_w, loc[1] * self.img_size[1] // img_h] for loc in seg_label["points"]
            ]
            index = self.label_index[label]  # 标签名称对应的标签数
            pts = np.array(points, dtype='int32').reshape((-1, 1, 2))  # 蒙板包围框
            mask = seg_mask[index]  # 蒙板
            cv.fillPoly(mask, [pts], color=1)  # 将蒙板画入分割标签

        origin_img, seg_mask = self.process(origin_img, seg_mask)
        img = cv.resize(origin_img, self.img_size, interpolation=cv.INTER_AREA)
        img = np.transpose(img, (2, 0, 1))  # 转化维度
        img = img[:, :] / 255  # 归一化
        return img, seg_mask

    def __len__(self):
        return len(self.datalist)
